import sys
import cv2
from PySide6.QtWidgets import QApplication, QWidget
from PySide6.QtGui import QImage, QPixmap
from PySide6.QtCore import QTimer, Qt  # 添加 Qt 导入
from ultralytics import YOLO
from ui_form import Ui_Widget

class VideoWindow(QWidget):
    def __init__(self):
        super().__init__()
        # 初始化 UI
        self.ui = Ui_Widget()
        self.ui.setupUi(self)

        # 加载 YOLOv8 模型
        self.model = YOLO("yolov8n.pt")

        # 初始化视频捕获（0 表示默认摄像头）
        self.cap = cv2.VideoCapture(0)
        self.is_running = False

        # 设置定时器，用于刷新视频帧
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.update_frame)

        # 连接按钮信号
        self.ui.startButton.clicked.connect(self.start_video)
        self.ui.stopButton.clicked.connect(self.stop_video)

    def start_video(self):
        if not self.is_running:
            if not self.cap.isOpened():
                self.cap = cv2.VideoCapture(0)
            self.is_running = True
            self.timer.start(30)

    def stop_video(self):
        if self.is_running:
            self.is_running = False
            self.timer.stop()
            self.cap.release()
            self.ui.videoLabel.clear()

    def update_frame(self):
        if self.is_running and self.cap.isOpened():
            ret, frame = self.cap.read()
            if ret:
                # 使用 YOLOv8 进行目标检测
                results = self.model(frame)
                annotated_frame = results[0].plot()

                # 将 OpenCV 图像（BGR）转换为 Qt 可用的格式（RGB）
                frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
                h, w, ch = frame_rgb.shape
                bytes_per_line = ch * w
                qt_image = QImage(frame_rgb.data, w, h, bytes_per_line, QImage.Format_RGB888)
                pixmap = QPixmap.fromImage(qt_image)

                # 在 QLabel 上显示图像，保持宽高比
                self.ui.videoLabel.setPixmap(pixmap.scaled(self.ui.videoLabel.size(), aspectMode=Qt.KeepAspectRatio))

    def closeEvent(self, event):
        # 窗口关闭时释放资源
        self.stop_video()
        event.accept()

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = VideoWindow()
    window.show()
    sys.exit(app.exec())